import numpy as np
from mpl_toolkits import mplot3d
from matplotlib import pyplot as plt
import tensorflow as tf
from python_ai.common.xcommon import *
import os

np.random.seed(777)
tf.random.set_seed(777)

# parameters
iters = 1000
alpha = 0.1

# (2)	求f(x) = a*x**2 + b*x + c的最小值
# ①	合理创建变量和常量
x = tf.Variable(0, dtype=tf.float32, name='x')
a = tf.constant(2.0, dtype=tf.float32, name='a')
b = tf.constant(3.0, dtype=tf.float32, name='b')
c = tf.constant(-4.0, dtype=tf.float32, name='c')
xmin = - b / (2. * a)

# ②	设定优化模型
def x_step():
    with tf.GradientTape() as tape:
        y = a * x ** 2 + b * x + c
    dydx = tape.gradient(y, x)
    tf.print(dydx)
    x.assign(x - alpha * dydx)

# ③	循环迭代计算
last_x = x.numpy()
for i in range(iters):
    x_step()
    tf.print(f'#{i + 1}: {x}', output_stream=sys.stdout)
    this_x = x.numpy()
    if np.isclose(last_x, this_x):
        print('Converged!')
        break
    last_x = this_x

# ④	打印相关数据值
print(f'Expected min: {xmin}')
print(f'Obtained min: {this_x}')
